from langgraph.graph import StateGraph, END
from states import GraphState
from langchain.prompts import PromptTemplate
from parsers import SQLOutputParser
from regexers import add_comments_to_top
from sql import get_database_metadata
from langchain_core.output_parsers import JsonOutputParser
# from loggings import log_ai_function
from llm import llm
from db import connection

function_name="sql_database_pipeline"
AGENT_NAME = "sql_database_agent"

sql_query_code_prompt_str = """
You are a SQL Database Coding Expert. Given the following information about the SQL database, 
write the SQL code to collect the data and process it according to user instructions. 
The code should be tailored to the SQL database characteristics and should take into account user instructions, recommended steps, database and table characteristics.

IMPORTANT INSTRUCTIONS:
- Do not use a LIMIT clause unless a user specifies a limit to be returned.
- Return SQL in ```sql ``` format.
- Only return a single query if possible.
- Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.
- Pay attention to the SQL dialect from the database summary metadata. Write the SQL code according to the dialect specified.
- IMPORTANT: Pay attention to the table names and column names in the database. Make sure to use the correct table and column names in the SQL code. If a space is present in the table name or column name, make sure to account for it.


User instructions / Question:
{user_instructions}

Recommended Steps:
{recommended_steps}

Below are summaries of the database metadata and the SQL tables:
{all_sql_database_summary}

Return:
- The SQL code in ```sql ``` format to collect the data and process it according to the user instructions.

Avoid these:
- Do not include steps to save files.
- Do not include steps to modify existing tables, create new tables or modify the database schema.
- Make sure not to alter the existing data in the database.
- Make sure not to include unsafe code that could cause data loss or corruption.

"""

def smart_schema_filter(llm, user_instructions, all_sql_database_summary, smart_filtering = True):
    """
    This function filters the tables and columns based on the user instructions and the recommended steps.
    """
    # Smart schema filtering
    if smart_filtering:
        print("    * SMART FILTER SCHEMA")
        
        filter_schema_prompt = PromptTemplate(
            template="""
            You are a highly skilled data engineer. The user question is:

                "{user_instructions}"

            You have the full database metadata in JSON format below:

                {all_sql_database_summary}
            
            
            Please return ONLY the subset of this metadata that is relevant to answering the user’s question. 
            - Preserve the same JSON structure for "schemas" -> "tables" -> "columns". 
            - If any schemas/tables are irrelevant, omit them entirely.
            - If some columns in a relevant table are not needed, you can still keep them if you aren't sure. 
            - However, try to keep only the minimum amount of data required to answer the user’s question.

            Return a valid JSON object. Do not include any additional explanation or text outside of the JSON.
            """,
            input_variables=["user_instructions", "full_metadata_json"]
        )
        
        filter_schema_agent = filter_schema_prompt | llm | JsonOutputParser()
        
        response = filter_schema_agent.invoke({
            "user_instructions": user_instructions,
            "all_sql_database_summary": all_sql_database_summary
        })
        
        return response
    else:
        return all_sql_database_summary
        
        
def create_sql_query_code(state: GraphState):

    all_sql_database_summary = get_database_metadata(connection, n_samples=1)
    print(f"    * all_sql_database_summary: {all_sql_database_summary}")
    print("    * SMART FILTER SCHEMA")
    all_sql_database_summary = smart_schema_filter(
        llm, 
        state.get("user_instructions"), 
        all_sql_database_summary, 
        smart_filtering=True
    )

    print("    * CREATE SQL QUERY CODE")
    print(f"    * all_sql_database_summary: {all_sql_database_summary}")

    # Prompt to get the SQL code from the LLM
    sql_query_code_prompt = PromptTemplate(
        template=sql_query_code_prompt_str,
        input_variables=["user_instructions", "recommended_steps", "all_sql_database_summary"]
    )
    
    sql_query_code_agent = sql_query_code_prompt | llm | SQLOutputParser()
    
    sql_query_code = sql_query_code_agent.invoke({
        "user_instructions": state.get("user_instructions"),
        "recommended_steps": state.get("recommended_steps"),
        "all_sql_database_summary": all_sql_database_summary
    })
    
    print(f"    * sql_query_code: {sql_query_code}")

    print("    * CREATE PYTHON FUNCTION TO RUN SQL CODE")
    
    # code format is important here, so we use f-string formatting to preserve the code structure.
    response = f"""
def {function_name}(connection):
    import pandas as pd
    import sqlalchemy as sql

    # Create a connection if needed
    is_engine = isinstance(connection, sql.engine.base.Engine)
    conn = connection.connect() if is_engine else connection

    sql_query = '''
    {sql_query_code}
    '''

    return pd.read_sql(sql_query, connection)
    """
    
    response = add_comments_to_top(response, AGENT_NAME)
        
    # For logging: store the code generated
    # file_path, file_name_2 = log_ai_function(
    #     response=response,
    #     file_name=file_name,
    #     log=log,
    #     log_path=log_path,
    #     overwrite=overwrite
    # )
    
    return {
        "sql_query_code": sql_query_code,
        "sql_database_function": response,
        # "sql_database_function_path": file_path,
        # "sql_database_function_file_name": file_name_2,
        "sql_database_function_name": function_name,
        "all_sql_database_summary": all_sql_database_summary
    }
